#!/usr/bin/env python3
"""
CasADi-based automatic differentiation model for quadrotor flatness mapping.

This script mirrors the computations in `quad_flatness_wang.hpp` using CasADi
symbolics so that the forward map, its Jacobian, and reverse-mode sensitivities
can be queried or code-generated. The original C++ implementation exposes a
`forward` routine that maps differential-flat outputs to attitude, thrust, and
body rates. This module preserves the same inputs and outputs while providing
automatic derivatives through CasADi.
"""

from __future__ import annotations

import casadi as ca


def normalize_with_grad(x: ca.MX, xd: ca.MX) -> tuple[ca.MX, ca.MX]:
    """Normalize `x` and propagate derivatives as in `normalizeWithGrad`."""
    x_sqr_norm = ca.dot(x, x)
    x_norm = ca.sqrt(x_sqr_norm)

    x_normalized = x / x_norm
    projection = ca.dot(x, xd) / x_sqr_norm
    x_normalized_dot = (xd - x * projection) / x_norm

    return x_normalized, x_normalized_dot


def rotation_to_quaternion(rot_m: ca.MX) -> ca.MX:
    """
    Convert a rotation matrix to a quaternion (w, x, y, z) following Eigen's
    convention. Branching is implemented with `ca.if_else` to keep symbolic
    differentiability.
    """
    trace = rot_m[0, 0] + rot_m[1, 1] + rot_m[2, 2]

    def case_trace_positive() -> ca.MX:
        s = ca.sqrt(trace + 1.0) * 2.0
        qw = 0.25 * s
        qx = (rot_m[2, 1] - rot_m[1, 2]) / s
        qy = (rot_m[0, 2] - rot_m[2, 0]) / s
        qz = (rot_m[1, 0] - rot_m[0, 1]) / s
        return ca.vertcat(qw, qx, qy, qz)

    def case_x_largest() -> ca.MX:
        s = ca.sqrt(1.0 + rot_m[0, 0] - rot_m[1, 1] - rot_m[2, 2]) * 2.0
        qw = (rot_m[2, 1] - rot_m[1, 2]) / s
        qx = 0.25 * s
        qy = (rot_m[0, 1] + rot_m[1, 0]) / s
        qz = (rot_m[0, 2] + rot_m[2, 0]) / s
        return ca.vertcat(qw, qx, qy, qz)

    def case_y_largest() -> ca.MX:
        s = ca.sqrt(1.0 + rot_m[1, 1] - rot_m[0, 0] - rot_m[2, 2]) * 2.0
        qw = (rot_m[0, 2] - rot_m[2, 0]) / s
        qx = (rot_m[0, 1] + rot_m[1, 0]) / s
        qy = 0.25 * s
        qz = (rot_m[1, 2] + rot_m[2, 1]) / s
        return ca.vertcat(qw, qx, qy, qz)

    def case_z_largest() -> ca.MX:
        s = ca.sqrt(1.0 + rot_m[2, 2] - rot_m[0, 0] - rot_m[1, 1]) * 2.0
        qw = (rot_m[1, 0] - rot_m[0, 1]) / s
        qx = (rot_m[0, 2] + rot_m[2, 0]) / s
        qy = (rot_m[1, 2] + rot_m[2, 1]) / s
        qz = 0.25 * s
        return ca.vertcat(qw, qx, qy, qz)

    x_largest = ca.logic_and(rot_m[0, 0] > rot_m[1, 1], rot_m[0, 0] > rot_m[2, 2])
    y_largest = rot_m[1, 1] > rot_m[2, 2]

    quat = ca.if_else(
        trace > 0.0,
        case_trace_positive(),
        ca.if_else(
            x_largest,
            case_x_largest(),
            ca.if_else(y_largest, case_y_largest(), case_z_largest()),
        ),
    )

    return quat


def build_flatness_forward_functions() -> dict[str, ca.Function]:
    """
    Construct CasADi functions for the flatness forward map and its Jacobian.

    Returns a dictionary containing:
        - forward: flatness map (thr, quat, omg)
        - jacobian: derivative of the stacked outputs with respect to inputs
        - backward: reverse-mode accumulation given output gradients
        - output_vector: symbolic output vector (thr followed by quat and omg)
    """
    vel = ca.MX.sym("vel", 3)
    acc = ca.MX.sym("acc", 3)
    jer = ca.MX.sym("jer", 3)
    psi = ca.MX.sym("psi")
    dpsi = ca.MX.sym("dpsi")

    zb, zbd = normalize_with_grad(acc, jer)

    syaw = ca.sin(psi)
    cyaw = ca.cos(psi)

    xc = ca.vertcat(cyaw, syaw, 0.0)
    xcd = ca.vertcat(-syaw * dpsi, cyaw * dpsi, 0.0)

    yc = ca.cross(zb, xc)
    ycd = ca.cross(zbd, xc) + ca.cross(zb, xcd)

    yb, ybd = normalize_with_grad(yc, ycd)

    xb = ca.cross(yb, zb)
    xbd = ca.cross(ybd, zb) + ca.cross(yb, zbd)

    rot_m = ca.hcat([xb, yb, zb])
    quat = rotation_to_quaternion(rot_m)

    # The original C++ routine does not assign `thr`.
    # We adopt the physically consistent choice ||acc|| as thrust magnitude.
    thr = ca.norm_2(acc)

    omg = ca.vertcat(
        (ca.dot(zb, ybd) - ca.dot(yb, zbd)) / 2.0,
        (ca.dot(xb, zbd) - ca.dot(zb, xbd)) / 2.0,
        (ca.dot(yb, xbd) - ca.dot(xb, ybd)) / 2.0,
    )

    outputs = ca.vertcat(thr, quat, omg)
    inputs = ca.vertcat(vel, acc, jer, psi, dpsi)

    forward = ca.Function(
        "flatness_forward",
        [vel, acc, jer, psi, dpsi],
        [thr, quat, omg],
        ["vel", "acc", "jer", "psi", "dpsi"],
        ["thr", "quat", "omg"],
    )

    jacobian = ca.jacobian(outputs, inputs)

    jacobian_fun = ca.Function(
        "flatness_forward_jacobian",
        [vel, acc, jer, psi, dpsi],
        [jacobian],
        ["vel", "acc", "jer", "psi", "dpsi"],
        ["doutput_dinput"],
    )

    thr_grad = ca.MX.sym("thr_grad")
    quat_grad = ca.MX.sym("quat_grad", 4)
    omg_grad = ca.MX.sym("omg_grad", 3)

    output_grads = ca.vertcat(thr_grad, quat_grad, omg_grad)
    input_grads = ca.mtimes(jacobian.T, output_grads)

    vel_grad = input_grads[0:3]
    acc_grad = input_grads[3:6]
    jer_grad = input_grads[6:9]
    psi_grad = input_grads[9]
    dpsi_grad = input_grads[10]

    backward = ca.Function(
        "flatness_backward",
        [vel, acc, jer, psi, dpsi, thr_grad, quat_grad, omg_grad],
        [vel_grad, acc_grad, jer_grad, psi_grad, dpsi_grad],
        ["vel", "acc", "jer", "psi", "dpsi", "thr_grad", "quat_grad", "omg_grad"],
        ["vel_grad", "acc_grad", "jer_grad", "psi_grad", "dpsi_grad"],
    )

    return {
        "forward": forward,
        "jacobian": jacobian_fun,
        "backward": backward,
        "output_vector": outputs,
    }


if __name__ == "__main__":
    functions = build_flatness_forward_functions()

    sample_vel = [1.0, 0.5, -0.2]
    sample_acc = [0.0, 0.0, 9.81]
    sample_jer = [0.1, -0.1, 0.05]
    sample_psi = 0.3
    sample_dpsi = 0.05

    thr_val, quat_val, omg_val = functions["forward"](
        sample_vel, sample_acc, sample_jer, sample_psi, sample_dpsi
    )
    jac_val = functions["jacobian"](
        sample_vel, sample_acc, sample_jer, sample_psi, sample_dpsi
    )
    vel_grad_val, acc_grad_val, jer_grad_val, psi_grad_val, dpsi_grad_val = functions[
        "backward"
    ](
        sample_vel,
        sample_acc,
        sample_jer,
        sample_psi,
        sample_dpsi,
        1.0,
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 0.0],
    )

    print("Thrust:", float(thr_val))
    print("Quaternion (w, x, y, z):", quat_val.T)
    print("Body rates:", omg_val.T)
    print("Jacobian shape:", jac_val.shape)
    print("Backward vel grad:", vel_grad_val.T)
    print("Backward acc grad:", acc_grad_val.T)
    print("Backward psi grad:", float(psi_grad_val))
